import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle

# 设置参数
x1, x2 = 0, 2  # 初始位移的区间
u0 = 1.0        # 最大位移值
a = 0.5         # 波速
t_max = 4.0     # 最大时间
num_frames = 100  # 动画帧数

# 计算中点
x0 = (x1 + x2) / 2

# 定义初始位移函数 φ(x)
def phi(x):
    """三角形初始位移函数"""
    result = np.zeros_like(x)
    
    # 第一段：x1 到中点 (x0)
    mask1 = (x >= x1) & (x <= x0)
    result[mask1] = 2 * u0 * (x[mask1] - x1) / (x2 - x1)
    
    # 第二段：中点 (x0) 到 x2
    mask2 = (x > x0) & (x <= x2)
    result[mask2] = 2 * u0 * (x2 - x[mask2]) / (x2 - x1)
    
    return result

# 定义波动方程的解 u(x, t)
def u(x, t):
    """达朗贝尔公式解"""
    return 0.5 * (phi(x + a*t) + phi(x - a*t))

# 创建空间网格
x_min, x_max = x1 - 2, x2 + 2  # 扩大显示范围
x = np.linspace(x_min, x_max, 1000)

# 设置图形
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用于显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正确显示负号

fig, ax = plt.subplots(figsize=(12, 8), facecolor='#f0f0f0')
fig.suptitle('达朗贝尔波动方程解', fontsize=16, fontweight='bold', color='#333333')

# 设置坐标轴
ax.set_xlim(x_min, x_max)
ax.set_ylim(-1.2*u0, 1.2*u0)
ax.set_xlabel('位置 (x)', fontsize=12)
ax.set_ylabel('位移 (u)', fontsize=12)
ax.grid(True, linestyle='--', alpha=0.7)
ax.set_facecolor('#f8f8f8')

# 绘制初始位移
initial_wave, = ax.plot(x, phi(x), 'b-', linewidth=2.5, alpha=0.7, label='初始位移 φ(x)')

# 绘制传播波
left_wave, = ax.plot([], [], 'g--', linewidth=1.8, alpha=0.8, label='向左传播波 φ(x+at)')
right_wave, = ax.plot([], [], 'r--', linewidth=1.8, alpha=0.8, label='向右传播波 φ(x-at)')
combined_wave, = ax.plot([], [], 'm-', linewidth=3, label='合成波 u(x,t)')

# 绘制时间标记
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=12, 
                    bbox=dict(facecolor='white', alpha=0.8))

# 添加图例
ax.legend(loc='upper right', fontsize=10)

# 添加说明
info_text = ax.text(0.02, 0.85, 
                   '达朗贝尔公式: u(x,t) = ½[φ(x+at) + φ(x-at)]\n'
                   f'波速 a = {a}, 最大位移 u0 = {u0}',
                   transform=ax.transAxes, fontsize=10, 
                   bbox=dict(facecolor='white', alpha=0.7))

# 添加关键点标记
ax.plot([x1, x2], [0, 0], 'ko', markersize=4)
ax.annotate(f'x1={x1}', (x1, -0.1), textcoords="offset points", xytext=(0,5), ha='center')
ax.annotate(f'x2={x2}', (x2, -0.1), textcoords="offset points", xytext=(0,5), ha='center')
ax.annotate(f'中点 x0={x0}', (x0, 0.1), textcoords="offset points", xytext=(0,5), ha='center')

# 初始化函数
def init():
    left_wave.set_data([], [])
    right_wave.set_data([], [])
    combined_wave.set_data([], [])
    time_text.set_text('时间 t = 0.0')
    return left_wave, right_wave, combined_wave, time_text

# 动画更新函数
def update(frame):
    t = frame * t_max / num_frames
    
    # 计算各波形
    left = phi(x + a*t)
    right = phi(x - a*t)
    total = u(x, t)
    
    # 更新波形数据
    left_wave.set_data(x, left)
    right_wave.set_data(x, right)
    combined_wave.set_data(x, total)
    
    # 更新时间文本
    time_text.set_text(f'时间 t = {t:.2f}')
    
    return left_wave, right_wave, combined_wave, time_text

# 创建动画
ani = FuncAnimation(fig, update, frames=num_frames,
                   init_func=init, blit=True, interval=50)

# 保存为GIF
ani.save('./dAlembert_wave_animation.gif', writer='pillow', fps=20, dpi=100)

# 显示动画
plt.tight_layout()
plt.show()